﻿// Inner Fire 游戏引擎库
// 场景分装 - 房间
//
// Copyright (c) 2023 by 尔西大帝. All rights reserved.
//
// 这个类分装各种场景的构建（都在一个文件里面太大了）
//
// Author: 尔西大帝
// Email: 2670613871@qq.com
// Created: 2024-11-24

#include <dx_graphics.h>

namespace ifire {
void DxGraphics::VecAdd_InitPipeline() {
  VecAdd_BuildBuffers();
  VecAdd_BuildRootSignature();
  VecAdd_BuildDescriptorHeaps();
  VecAdd_BuildShadersAndInputLayout();
  VecAdd_BuildFrameResources();
  VecAdd_BuildPSOs();
}

// 这东西就是什么都不做而已！
void DxGraphics::VecAdd_Draw() {}

void DxGraphics::VecAdd_BuildBuffers() {
  // 生成一些数据
  std::vector<Data> dataA(NUM_DATA_ELEMETS);
  std::vector<Data> dataB(NUM_DATA_ELEMETS);
  for (int i = 0; i < NUM_DATA_ELEMETS; ++i) {
    float fi = static_cast<float>(i);
    dataA[i].v1 = XMFLOAT3(fi, fi, fi);
    dataA[i].v2 = XMFLOAT2(fi, 0);

    dataB[i].v1 = XMFLOAT3(-fi, fi, 0.0f);
    dataB[i].v2 = XMFLOAT2(0, -fi);
  }

  UINT64 byteSize = dataA.size() * sizeof(Data);

  // 创建SRV的BUFFER用做使用.
  mInputBufferA =
      CreateDefaultBuffer(dataA.data(), byteSize, mInputUploadBufferA);

  mInputBufferB =
      CreateDefaultBuffer(dataB.data(), byteSize, mInputUploadBufferB);

  // 创建UAV
  auto default_props = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_DEFAULT);
  auto res_desc = CD3DX12_RESOURCE_DESC::Buffer(
      byteSize, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS);
  HASSERT(device_->CreateCommittedResource(&default_props, D3D12_HEAP_FLAG_NONE,
      &res_desc, D3D12_RESOURCE_STATE_UNORDERED_ACCESS, nullptr,
      IID_PPV_ARGS(&mOutputBuffer)));

  // 注意这个堆的类型，是D3D12_HEAP_TYPE_READBACK。
  auto readback_props = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_READBACK);
  auto res_desc2 = CD3DX12_RESOURCE_DESC::Buffer(byteSize);
  HASSERT(device_->CreateCommittedResource(&readback_props,
      D3D12_HEAP_FLAG_NONE, &res_desc2, D3D12_RESOURCE_STATE_COPY_DEST, nullptr,
      IID_PPV_ARGS(&mReadBackBuffer)));
}

void DxGraphics::VecAdd_BuildRootSignature() {
  // 根参数可以是描述符表，根描述符，或根常量
  CD3DX12_ROOT_PARAMETER slot_root_paramater[3];
  // 这个是两个输入纹理
  slot_root_paramater[0].InitAsShaderResourceView(0);
  slot_root_paramater[1].InitAsShaderResourceView(1);
  // 这个绑到UAV0
  slot_root_paramater[2].InitAsUnorderedAccessView(0);

  // 根签名由一组根参数构成
  // 注意这里和之前的变化，增加了采样器的数量和地址。
  CD3DX12_ROOT_SIGNATURE_DESC root_signature_desc(
      3, slot_root_paramater, 0, nullptr, D3D12_ROOT_SIGNATURE_FLAG_NONE);

  // 序列化根签名
  ComPtr<ID3DBlob> serialized_root_sign = nullptr;
  ComPtr<ID3DBlob> error_blob = nullptr;
  HRESULT hr = D3D12SerializeRootSignature(&root_signature_desc,
      D3D_ROOT_SIGNATURE_VERSION_1, serialized_root_sign.GetAddressOf(),
      error_blob.GetAddressOf());
  // 常规错误检查
  if (error_blob != nullptr) {
    dxutils::log((char*)error_blob->GetBufferPointer());
  }
  HASSERT(hr);

  // 创建根签名
  HASSERT(device_->CreateRootSignature(0, //
      serialized_root_sign->GetBufferPointer(),
      serialized_root_sign->GetBufferSize(),
      IID_PPV_ARGS(root_signature_.GetAddressOf())));
}

void DxGraphics::VecAdd_BuildDescriptorHeaps() {}
void DxGraphics::VecAdd_BuildShadersAndInputLayout() {
  shaders_["VecAddCS"] =
      dxutils::compile_shader("Shaders\\VecAdd.hlsl", nullptr, "CS", "cs_5_0");
}
void DxGraphics::VecAdd_BuildFrameResources() {
  FrameResourceDesc desc = {};
  desc.PassCount = 1;

  for (int i = 0; i < FRAME_RESOURCE_COUNT; i++) {
    frame_resources_.push_back(
        std::make_unique<FrameResource>(device_.Get(), &desc));
  }
}
void DxGraphics::VecAdd_BuildPSOs() {
  D3D12_COMPUTE_PIPELINE_STATE_DESC computePsoDesc = {};
  computePsoDesc.pRootSignature = root_signature_.Get();
  computePsoDesc.CS = {
      reinterpret_cast<BYTE*>(shaders_["VecAddCS"]->GetBufferPointer()),
      shaders_["VecAddCS"]->GetBufferSize()};
  computePsoDesc.Flags = D3D12_PIPELINE_STATE_FLAG_NONE;
  HASSERT(device_->CreateComputePipelineState(
      &computePsoDesc, IID_PPV_ARGS(&pipelines_["VecAdd"])));
}

/// <summary>
/// 这个东西的主要流程其实和绘制一帧没有区别。在初始化一切之后，可以进行一帧的绘制。而不是在Draw循环中。
/// </summary>
void DxGraphics::VecAdd_DoComputeWork() {
  auto& cmd = command_list_;

  // 常规的帧初始化处理
  ThrowIfFailed(command_allocator_->Reset());
  ThrowIfFailed(
      cmd->Reset(command_allocator_.Get(), pipelines_["VecAdd"].Get()));
  cmd->SetComputeRootSignature(root_signature_.Get());

  // 将输出和输出项目填入（我终于明白DirectX为什么叫RootSignature这么奇葩的名称了）
  // 因为就和调用函数填参数一样。倒也不是非常“反人类”。
  cmd->SetComputeRootShaderResourceView(
      0, mInputBufferA->GetGPUVirtualAddress());
  cmd->SetComputeRootShaderResourceView(
      1, mInputBufferB->GetGPUVirtualAddress());
  cmd->SetComputeRootUnorderedAccessView(
      2, mOutputBuffer->GetGPUVirtualAddress());

  // 就是这行，进行了计算着色器。
  // 计算着色器能够在任何时候进行，就是这个东西Dispatch。然后更改后的数据在各种BUFFER中。
  // 然后接下来就是倒腾各种BUFFER里面的数据，传去这，用到那儿而已。
  cmd->Dispatch(1, 1, 1);

  // 这个就是把OutputBuffer的东西倒腾出来到ReadBackBuffer。然后就可以用CPU读取。
  auto barrier = CD3DX12_RESOURCE_BARRIER::Transition(mOutputBuffer.Get(),
      D3D12_RESOURCE_STATE_COMMON, D3D12_RESOURCE_STATE_COPY_SOURCE);
  cmd->ResourceBarrier(1, &barrier);

  cmd->CopyResource(mReadBackBuffer.Get(), mOutputBuffer.Get());

  auto barrier2 = CD3DX12_RESOURCE_BARRIER::Transition(mOutputBuffer.Get(),
      D3D12_RESOURCE_STATE_COPY_SOURCE, D3D12_RESOURCE_STATE_COMMON);
  cmd->ResourceBarrier(1, &barrier2);

  // Done recording commands.
  ThrowIfFailed(cmd->Close());

  // Add the command list to the queue for execution.
  ID3D12CommandList* cmdsLists[] = {cmd.Get()};
  command_queue_->ExecuteCommandLists(_countof(cmdsLists), cmdsLists);

  // Wait for the work to finish.
  FlushCommandQueue();

  // Map the data so we can read it on CPU.
  Data* mappedData = nullptr;
  ThrowIfFailed(
      mReadBackBuffer->Map(0, nullptr, reinterpret_cast<void**>(&mappedData)));

  std::ofstream fout("results.txt");

  for (int i = 0; i < NUM_DATA_ELEMETS; ++i) {
    fout << "(" << mappedData[i].v1.x << ", " << mappedData[i].v1.y << ", "
         << mappedData[i].v1.z << ", " << mappedData[i].v2.x << ", "
         << mappedData[i].v2.y << ")" << std::endl;
  }

  mReadBackBuffer->Unmap(0, nullptr);

  dxutils::log("计算结果已经写入：results.txt");
}

} // namespace ifire